# -*- coding: utf-8 -*-
"""
车牌字符识别模块

基于自定义CNN网络的车牌字符识别功能，包括：
- 车牌图像预处理
- 字符序列解码（贪心去重）
- 置信度计算
- 模型初始化和推理

Author: AI Assistant
Date: 2025-06-12
"""

import torch
import torch.nn as nn
import cv2
import numpy as np
import os
import time
import sys

from plate_recognition.plateNet import myNet_ocr

# 中国车牌字符集定义
PLATE_CHARACTERS = "#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航深0123456789ABCDEFGHJKLMNPQRSTUVWXYZ"

# 图像标准化参数
IMAGE_MEAN_VALUE = 0.588
IMAGE_STD_VALUE = 0.193

def read_image_with_chinese_path(image_path):
    """
    读取包含中文路径的图像文件
    
    Args:
        image_path: 图像文件路径
        
    Returns:
        OpenCV格式的图像数组
    """
    image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), -1)
    return image

def get_all_image_files(root_directory, file_list):
    """
    递归获取目录下所有图像文件
    
    Args:
        root_directory: 根目录路径
        file_list: 用于存储文件路径的列表
    """
    try:
        files = os.listdir(root_directory)
        for file_name in files:
            file_path = os.path.join(root_directory, file_name)
            
            if os.path.isfile(file_path):
                # 检查是否为图像文件
                image_extensions = ['.jpg', '.png', '.JPG', '.bmp', '.jpeg', '.JPEG']
                if any(file_name.endswith(ext) for ext in image_extensions):
                    file_list.append(file_path)
            else:
                # 递归处理子目录
                get_all_image_files(file_path, file_list)
    except PermissionError:
        print(f"权限不足，无法访问目录: {root_directory}")

def decode_character_sequence(predictions, confidence_scores):
    """
    解码车牌字符序列，去除重复字符并计算置信度
    
    Args:
        predictions: 字符预测结果数组
        confidence_scores: 字符置信度分数数组
        
    Returns:
        tuple: (去重后的字符索引列表, 整体置信度分数, 各字符置信度列表)
    """
    previous_char = 0
    decoded_characters = []
    overall_confidence = 1.0
    character_confidences = []
    selected_confidences = []
    
    for i, current_char in enumerate(predictions):
        if current_char != 0:  # 非空白字符
            if current_char != previous_char:
                # 新字符
                decoded_characters.append(current_char)
                selected_confidences.append(confidence_scores[i])
            elif confidence_scores[i] > confidence_scores[i-1]:
                # 相同字符但置信度更高，替换之前的
                selected_confidences.pop()
                selected_confidences.append(confidence_scores[i])
        previous_char = current_char
    
    # 计算整体置信度和各字符置信度
    for confidence in selected_confidences:
        overall_confidence *= confidence
        char_confidence = int(round(confidence * 100))
        character_confidences.append(char_confidence)
    
    overall_confidence = int(round(overall_confidence * 100))
    
    return decoded_characters, overall_confidence, character_confidences

def preprocess_plate_image(image, device):
    """
    预处理车牌图像用于模型推理
    
    Args:
        image: 输入的车牌图像
        device: 计算设备(CPU/GPU)
        
    Returns:
        预处理后的张量
    """
    # 调整图像尺寸为模型要求的大小
    resized_image = cv2.resize(image, (168, 48))
    processed_image = np.reshape(resized_image, (48, 168, 3))

    # 图像标准化
    processed_image = processed_image.astype(np.float32)
    processed_image = (processed_image / 255.0 - IMAGE_MEAN_VALUE) / IMAGE_STD_VALUE
    
    # 转换为PyTorch张量格式 (CHW)
    processed_image = processed_image.transpose([2, 0, 1])
    image_tensor = torch.from_numpy(processed_image)

    # 移动到指定设备并添加批次维度
    image_tensor = image_tensor.to(device)
    image_tensor = image_tensor.view(1, *image_tensor.size())
    
    return image_tensor

def extract_confidence_scores(probability_matrix):
    """
    从概率矩阵中提取置信度分数
    
    Args:
        probability_matrix: 模型输出的概率矩阵
        
    Returns:
        置信度分数列表
    """
    # 转换为numpy数组
    probs = probability_matrix.cpu().detach().numpy()
    confidence_scores = []
    
    for batch_idx in range(probs.shape[0]):
        prob_batch = probs[batch_idx, :, :]
        
        for time_step in range(prob_batch.shape[1]):
            max_confidence = max(prob_batch[:, time_step])
            confidence_scores.append(max_confidence)
    
    return confidence_scores

def recognize_plate_characters(plate_image, device, model):
    """
    识别车牌图像中的字符
    
    Args:
        plate_image: 车牌图像
        device: 计算设备
        model: 预训练的字符识别模型
        
    Returns:
        tuple: (车牌字符串, 整体置信度, 各字符置信度列表)
    """
    # 图像预处理
    input_tensor = preprocess_plate_image(plate_image, device)
    
    # 模型推理
    with torch.no_grad():
        character_predictions, probability_scores = model(input_tensor)
    
    # 转换预测结果
    predictions = character_predictions.view(-1).detach().cpu().numpy()
    confidence_scores = extract_confidence_scores(probability_scores)
    
    # 解码字符序列
    decoded_chars, overall_confidence, char_confidences = decode_character_sequence(
        predictions, confidence_scores
    )
    
    # 转换为字符串
    plate_text = ""
    for char_index in decoded_chars:
        plate_text += PLATE_CHARACTERS[char_index]

    return plate_text, overall_confidence, char_confidences

def initialize_recognition_model(device, model_path):
    """
    初始化车牌字符识别模型
    
    Args:
        device: 计算设备
        model_path: 模型文件路径
        
    Returns:
        初始化后的模型
    """
    try:
        # 加载模型检查点
        checkpoint = torch.load(model_path, map_location=device, weights_only=False)
        model_state = checkpoint['state_dict']
        model_config = checkpoint['cfg']
        
        # 创建模型实例
        model = myNet_ocr(num_classes=78, export=True, cfg=model_config)
        
        # 加载权重
        model.load_state_dict(model_state)
        model.to(device)
        model.eval()
        
        return model
        
    except Exception as e:
        print(f"模型初始化失败: {e}")
        raise

# 向后兼容的函数别名
cv_imread = read_image_with_chinese_path
allFilePath = get_all_image_files  
decodePlate = decode_character_sequence
image_processing = preprocess_plate_image
lprnet_score_decode = extract_confidence_scores
get_plate_result = recognize_plate_characters
init_model = initialize_recognition_model

if __name__ == '__main__':
    """
    模块测试入口
    """
    # 测试配置
    test_image_path = "images/tmp2424.png"
    test_directory = "double_plate"
    
    # 初始化设备
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    print(f"使用设备: {device}")
    
    # 获取测试文件列表
    test_files = []
    if os.path.exists(test_directory):
        get_all_image_files(test_directory, test_files)
        print(f"找到 {len(test_files)} 个测试图像")
    else:
        print(f"测试目录不存在: {test_directory}")
    
    if not test_files:
        print("没有找到测试图像文件")
        sys.exit(1)
    
    # 初始化模型
    try:
        model_path = "plate_recognition/model/plate_rec.pth"
        recognition_model = initialize_recognition_model(device, model_path)
        print("模型初始化成功")
    except Exception as e:
        print(f"模型初始化失败: {e}")
        sys.exit(1)
    
    # 批量测试
    correct_predictions = 0
    start_time = time.time()
    
    for image_path in test_files:
        try:
            # 识别车牌
            predicted_plate, confidence, char_confidences = recognize_plate_characters(
                image_path, device, recognition_model
            )
            
            # 从文件名提取真实车牌号（假设文件名格式为"车牌号_xxx.jpg"）
            filename = os.path.basename(image_path)
            actual_plate = filename.split('_')[0]
            
            # 统计准确率
            if predicted_plate == actual_plate:
                correct_predictions += 1
            else:
                print(f"识别错误: {actual_plate} -> {predicted_plate} ({image_path})")
                
        except Exception as e:
            print(f"处理图像失败 {image_path}: {e}")
    
    end_time = time.time()
    
    # 输出测试结果
    total_images = len(test_files)
    accuracy = correct_predictions / total_images if total_images > 0 else 0
    elapsed_time = end_time - start_time
    
    print(f"\n测试结果:")
    print(f"总图像数: {total_images}")
    print(f"正确识别: {correct_predictions}")
    print(f"准确率: {accuracy:.4f}")
    print(f"总耗时: {elapsed_time:.2f}秒")
    print(f"平均耗时: {elapsed_time/total_images:.3f}秒/图像")
        
